{
 "cells": [
  {
   "cell_type": "raw",
   "id": "ef5b30a5-14f9-44cb-853b-a043b8e22c61",
   "metadata": {},
   "source": [
    "dataset/\n",
    "│\n",
    "├── images/\n",
    "│   ├── train/         # 训练集图片\n",
    "│   ├── val/           # 验证集图片\n",
    "|   ├── test/\n",
    "│\n",
    "├── labels/\n",
    "│   ├── train/         # 训练集对应的标注文件\n",
    "│   ├── val/           # 验证集对应的标注文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "908f4497-f56a-47f6-827b-e5434da0aac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import shutil\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "47eb68b0-0f8f-4f68-91f8-58625390de25",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mkdir():\n",
    "    img_train_save_folder = \"../datasets/kitti/images/train\"\n",
    "    img_val_save_folder = \"../datasets/kitti/images/val\"\n",
    "    label_train_save_folder = \"../datasets/kitti/labels/train\"\n",
    "    label_val_save_folder = \"../datasets/kitti/labels/val\"\n",
    "    try: \n",
    "        os.makedirs(img_train_save_folder)\n",
    "    except:\n",
    "        None\n",
    "    try: \n",
    "        os.makedirs(img_val_save_folder)\n",
    "    except:\n",
    "        None\n",
    "    try: \n",
    "        os.makedirs(label_train_save_folder)\n",
    "    except:\n",
    "        None\n",
    "    try: \n",
    "        os.makedirs(label_val_save_folder)\n",
    "    except:\n",
    "        None\n",
    "    return (img_train_save_folder, img_val_save_folder, label_train_save_folder, label_val_save_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5e6ad222-3985-4bc8-95c8-313dc8666837",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cp_2_target(img_path, label_path, img_save_folder, label_save_folder, img_name):\n",
    "    label_name = os.path.splitext(img_name)[0] + \".txt\"\n",
    "    train_img_path = os.path.join(img_path, img_name)\n",
    "    train_label_path = os.path.join(label_path, label_name)\n",
    "    \n",
    "    if not os.path.exists(train_img_path) or not os.path.exists(train_label_path):\n",
    "        sys.exit()\n",
    "\n",
    "    shutil.copy(train_img_path, img_save_folder)\n",
    "    shutil.copy(train_label_path, label_save_folder)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "03e4fa95-b716-4bdc-becd-32b8fe5bb595",
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_dataset(img_path, label_path = \"../converted_label/k_2_y\", shuffle_img = True, val_percent = 0.2):\n",
    "    img_train_save_folder, img_val_save_folder, label_train_save_folder, label_val_save_folder = mkdir()#先拿到保存路径\n",
    "\n",
    "    all_imgs = os.listdir(img_path) # 得到所有文件的名字\n",
    "    if shuffle_img:\n",
    "        #打乱所有图片\n",
    "        random.shuffle(all_imgs)\n",
    "\n",
    "    img_num = len(all_imgs)\n",
    "    train_num = int(img_num * (1 - val_percent))\n",
    "    val_num = img_num - train_num\n",
    "    train_imgs = all_imgs[ : train_num]\n",
    "    val_imgs = all_imgs[train_num : ]\n",
    "\n",
    "    #开始写入目标文件夹\n",
    "    for i, train_img in enumerate(train_imgs):\n",
    "        if train_img.endswith(('.jpg', '.png', '.jpeg')):\n",
    "            cp_2_target(img_path, \n",
    "                        label_path, \n",
    "                        img_train_save_folder, \n",
    "                        label_train_save_folder,  \n",
    "                        train_img) #得到一个图片和label的路径\n",
    "        if i % (train_num // 10) == 0:\n",
    "            progress = (i / train_num) * 100\n",
    "            print(f\"已完成 {progress:.0f}% ({i}/{train_num})\")\n",
    "            \n",
    "\n",
    "    for i, val_img in enumerate(val_imgs):\n",
    "        if val_img.endswith(('.jpg', '.png', '.jpeg')):\n",
    "            cp_2_target(img_path, \n",
    "                        label_path, \n",
    "                        img_val_save_folder, \n",
    "                        label_val_save_folder,  \n",
    "                        val_img) #得到一个图片和label的路径\n",
    "        if i % (val_num // 10) == 0:\n",
    "            progress = (i / val_num) * 100\n",
    "            print(f\"已完成 {progress:.0f}% ({i}/{val_num})\")\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5ff54974-934b-4bc9-801c-8c5589d8903e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "已完成 0% (0/4189)\n",
      "已完成 10% (418/4189)\n",
      "已完成 20% (836/4189)\n",
      "已完成 30% (1254/4189)\n",
      "已完成 40% (1672/4189)\n",
      "已完成 50% (2090/4189)\n",
      "已完成 60% (2508/4189)\n",
      "已完成 70% (2926/4189)\n",
      "已完成 80% (3344/4189)\n",
      "已完成 90% (3762/4189)\n",
      "已完成 100% (4180/4189)\n",
      "已完成 0% (0/1048)\n",
      "已完成 10% (104/1048)\n",
      "已完成 20% (208/1048)\n",
      "已完成 30% (312/1048)\n",
      "已完成 40% (416/1048)\n",
      "已完成 50% (520/1048)\n",
      "已完成 60% (624/1048)\n",
      "已完成 69% (728/1048)\n",
      "已完成 79% (832/1048)\n",
      "已完成 89% (936/1048)\n",
      "已完成 99% (1040/1048)\n"
     ]
    }
   ],
   "source": [
    "split_dataset(img_path = '/root/autodl-tmp/training/image_2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6eb140c-bbe7-4cea-9053-1f5e51576678",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "28fd3b25-ea92-4430-80ad-359a3422c359",
   "metadata": {},
   "outputs": [],
   "source": [
    "from MyTools import kitti_2_yolo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "330dce82-88ef-4f60-ad86-97f1c3c358aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 把测试集也弄到datasets里\n",
    "def cp_testsets(test_path = '/root/autodl-tmp/testing/image_2', test_labels_path = None):\n",
    "    img_test_save_folder = \"../datasets/kitti/images/test/\"\n",
    "    try: \n",
    "        shutil.rmtree(img_test_save_folder)\n",
    "    except:\n",
    "        None\n",
    "    os.makedirs(img_test_save_folder)\n",
    "\n",
    "    files = os.listdir(test_path)\n",
    "    for file in files:\n",
    "        file_path = os.path.join(test_path, file)\n",
    "        if os.path.isfile(file_path):\n",
    "            shutil.copy(file_path, img_test_save_folder)\n",
    "            \n",
    "    if test_labels_path is not None:\n",
    "        label_test_save_folder = \"../datasets/kitti/labels/test/\"\n",
    "        try: \n",
    "            shutil.rmtree(label_test_save_folder)\n",
    "        except:\n",
    "            None\n",
    "        os.makedirs(label_test_save_folder)\n",
    "        kitti_2_yolo(img_test_save_folder, test_labels_path, save_path = label_test_save_folder)\n",
    "        \n",
    "        \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "658b3e70-71f9-4cfd-a8d8-291052af1050",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1000\n",
      "2000\n"
     ]
    }
   ],
   "source": [
    "cp_testsets(test_labels_path = '/root/autodl-tmp/testing/label_2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd7beea-4607-4221-9b2b-57206ffeea45",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
